##############################
# STATE REACH AND DEVELOPMENT // PSRM // Replication
#
# Functions for voronoi cells vs. square grid cells comparison
# Sourced from vor_vs_cells.R
##############################

# (function) Get constitutive parts of admin units that are never split
get_adm_const_parts <- function(admin.shp, base.year = 1992, base.cowcode = NULL, ncore = 10, size.cutoff = .001){
  # Init
  require(snow)
  require(parallel)
  require(foreach)
  require(doParallel)
  
  # Cowcodes to work on
  base.cowcodes <- unique(admin.shp$cowcode[admin.shp$startyr <= base.year & admin.shp$endyr >= base.year])
  if(!is.null(base.cowcode)){
    base.cowcodes <- base.cowcode
  }
  
  if(ncore > 1){
    # Setup cluster
    print(paste(Sys.time(), "Setup cluster, load objects"))
    cl <- makeCluster(getOption("cl.cores", ncore))
    export.ls <- c("admin.shp")
    clusterExport(cl, export.ls,
                  envir = environment())
    clusterExport(cl, c("remove.holes"),
                  envir = .GlobalEnv)
    registerDoParallel(cl)
    
    # Compute constitutive parts
    const.parts <- foreach(c = base.cowcodes, .noexport = export.ls, 
                           .packages = c("rgdal", "rgeos", "raster", "sp"), 
                           .options.multicore = list(preschedule = FALSE)) %dopar% {
                             # Base shape
                             base.shape <- admin.shp[admin.shp$startyr <= base.year & admin.shp$endyr >= base.year & admin.shp$cowcode == c,]
                             base.shape <- gUnaryUnion(base.shape, id = rep(1, length(base.shape)))
                             
                             # Delete small holes
                             base.shape <- remove.holes(base.shape, max.size = 1)
                             
                             # Intersect with all shapes
                             adm.sub <- admin.shp[gIntersects(admin.shp,base.shape, byid = T)[1,],]
                             adm.sub <- raster::crop(adm.sub, base.shape)
                             
                             # Delete below size cutoff
                             adm.sub <- adm.sub[gArea(adm.sub, byid = T) > size.cutoff,]
                             
                             # As lines (with buffer)
                             adm.lines <- as(adm.sub, "SpatialLines")
                             adm.lines <- gBuffer(adm.lines, width = .00001)
                             
                             # Split base polygon
                             base.split <- gDifference(base.shape, adm.lines)     
                             
                             # Get holes
                             holes <- lapply(1:length(base.split@polygons[[1]]@Polygons), function(p){
                               if(base.split@polygons[[1]]@Polygons[[p]]@hole){
                                 Polygons(list(base.split@polygons[[1]]@Polygons[[p]]), ID = paste0(c, ".",p))
                               } else {
                                 NULL
                               }
                             })
                             holes[sapply(holes, is.null)] <- NULL
                             holes <- SpatialPolygons(holes)
                             
                             # Delete below size cutoff
                             holes <- holes[gArea(holes, byid = T) > size.cutoff,]
                             
                             # Valid polygons
                             base.split <- lapply(1:length(base.split@polygons[[1]]@Polygons), function(p){
                               if(base.split@polygons[[1]]@Polygons[[p]]@hole){
                                 NULL
                               } else {
                                 Polygons(list(base.split@polygons[[1]]@Polygons[[p]]), ID = paste0(c, ".",p))
                               }
                             })
                             base.split[sapply(base.split, is.null)] <- NULL
                             base.split <- SpatialPolygons(base.split)
                             
                             # Delete below size cutoff
                             base.split <- base.split[gArea(base.split, byid = T) > size.cutoff,]
                             
                             # Correct holes
                             if(length(holes) > 0){
                               save.rn <- row.names(base.split)
                               for(h in c(1:length(holes))){
                                 this.h <- holes[h, ]
                                 this.p.id <- which(gContains(base.split, this.h, byid = T))
                                 base.split@polygons[[this.p.id]] <- gDifference(base.split[this.p.id,], this.h)@polygons[[1]]    
                               }
                               row.names(base.split) <- save.rn
                             }
                             
                             # Return
                             SpatialPolygonsDataFrame(base.split, data.frame(const.id = 1:length(base.split),
                                                                             base.cowcode = c), match.ID = F)
                           }
    const.parts <- do.call(rbind, const.parts)
    
    # StopCluster
    stopCluster(cl)
    rm(cl)
  } else {
    const.parts <- lapply(base.cowcodes, function(c){
      # Base shape
      base.shape <- admin.shp[admin.shp$startyr <= base.year & admin.shp$endyr >= base.year & admin.shp$cowcode == c,]
      if(length(base.shape) == 0){return(NULL)}
      base.shape <- gUnaryUnion(base.shape, id = rep(1, length(base.shape)))
      
      # Delete small holes
      base.shape <- remove.holes(base.shape, max.size = 1)
      
      # Intersect with all shapes
      adm.sub <- admin.shp[gIntersects(admin.shp,base.shape, byid = T)[1,],]
      adm.sub <- raster::crop(adm.sub, base.shape)
      
      # Delete below size cutoff
      adm.sub <- adm.sub[gArea(adm.sub, byid = T) > size.cutoff,]
      
      # As lines (with buffer)
      adm.lines <- as(adm.sub, "SpatialLines")
      adm.lines <- gBuffer(adm.lines, width = .00001)
      adm.lines <- gUnaryUnion(adm.lines, id = rep(1, length(adm.lines)))
      
      # Split base polygon
      base.split <- gDifference(base.shape, adm.lines)     
      
      # Get holes
      holes <- lapply(1:length(base.split@polygons[[1]]@Polygons), function(p){
        if(base.split@polygons[[1]]@Polygons[[p]]@hole){
          Polygons(list(base.split@polygons[[1]]@Polygons[[p]]), ID = paste0(c, ".",p))
        } else {
          NULL
        }
      })
      holes[sapply(holes, is.null)] <- NULL
      holes <- SpatialPolygons(holes)
      
      # Delete below size cutoff
      holes <- holes[gArea(holes, byid = T) > size.cutoff,]
      
      # Valid polygons
      base.split <- lapply(1:length(base.split@polygons[[1]]@Polygons), function(p){
        if(base.split@polygons[[1]]@Polygons[[p]]@hole){
          NULL
        } else {
          Polygons(list(base.split@polygons[[1]]@Polygons[[p]]), ID = paste0(c, ".",p))
        }
      })
      base.split[sapply(base.split, is.null)] <- NULL
      base.split <- SpatialPolygons(base.split)
      
      # Delete below size cutoff
      base.split <- base.split[gArea(base.split, byid = T) > size.cutoff,]
      
      # Correct holes
      
      if(length(holes) > 0){
        save.rn <- row.names(base.split)
        for(h in c(1:length(holes))){
          this.h <- holes[h, ]
          this.p.id <- which(gContains(base.split, this.h, byid = T))
          base.split@polygons[[this.p.id]] <- gDifference(base.split[this.p.id,], this.h)@polygons[[1]]    
        }
        row.names(base.split) <- save.rn
      }
      
      # Return
      SpatialPolygonsDataFrame(base.split, data.frame(const.id = 1:length(base.split),
                                                      base.cowcode = c), match.ID = F)
    })
    const.parts <- do.call(rbind, const.parts)
  }
  
  
  # Return
  return(const.parts)
}

# SAMPLE VORONOI CELLS WITHIN AREA #####

# (function) Sample regular voronoi cells within entries of SpatialPolygonsDataFrame
sample_vorcells <- function(spdf, res, size, size.km2 = F, seed = 1, iter.max = 100, sample.type = "nonaligned",
                            ncore = 1){
  require(rgeos)
  if(ncore == 1){
    # Sampe cells per entry in spdf
    result <- do.call(rbind,
                      lapply(c(1:length(spdf)), function(i){
                        print(i)
                        t <- kmean_spat(shape = spdf[i,], res = res, size = size, 
                                        size.km2 = size.km2, seed = seed, iter.max = iter.max, sample.type = sample.type)
                        t$unit.id <- i
                        proj4string(t) <- proj4string(spdf)
                        row.names(t) <- paste0(i, ".", row.names(t))
                        return(t)
                      }))
  } else if(ncore > 1) {
    
    # Libraries
    require(foreach)
    require(doParallel)
    
    # Setup cluster
    try(registerDoSEQ())
    cl <- makeCluster(getOption("cl.cores", ncore), outfile = "error.txt")
    clusterExport(cl, list("voronoi_poly", "kmean_spat"), envir = .GlobalEnv)
    registerDoParallel(cl)
    
    # Compute tesselation
    result <- foreach(i = c(1:length(spdf)), .packages = c("deldir","raster","rgeos","geosphere","sp"),
                      .options.multicore = list(preschedule = FALSE)) %dopar% {
                        t <- kmean_spat(shape = spdf[i,], res = res, size = size, 
                                        size.km2 = size.km2, seed = seed, iter.max = iter.max, sample.type = sample.type)
                        t$unit.id <- i
                        proj4string(t) <- proj4string(spdf)
                        row.names(t) <- paste0(i, ".", row.names(t))
                        t
                      }
    result <- do.call(rbind, result)
    
    # Close cluster
    stopCluster(cl)
    rm(cl)
    
  } else {
    stop("Please enter a valid number of cores > 0")
  }
  
  # Clean result
  result <- gBuffer(result, byid = T, width = 0)
  
  # Return
  return(result)
}

# (function) Spatial Voronoi Polygons
voronoi_poly <- function(vpoints, extent = NULL){
  require(deldir)
  require(sp)
  
  # ... tesselation
  z = deldir(data.frame(vpoints@coords), rw = as.vector(extent))
  w = tile.list(z)
  
  # ... make polygons
  polys = vector(mode='list', length=length(w))
  for (i in seq(along=polys)) {
    pcrds = cbind(w[[i]]$x, w[[i]]$y)
    pcrds = rbind(pcrds, pcrds[1,])
    polys[[i]] = Polygons(list(Polygon(pcrds)), ID=as.character(w[[i]]$ptNum))
  }
  return(SpatialPolygonsDataFrame(SpatialPolygons(polys),
                                  data = data.frame(cell.id = z$ind.orig), 
                                  match.ID = F))
}

# (function) Use kmeans-clustering for optimized voronoi tesselation of a delimited plane
kmean_spat <- function(shape, res, size, size.km2 = F,
                       seed = 1, iter.max = 100, sample.type = "nonaligned"){
  # Packages
  require(sp)
  require(deldir)
  require(geosphere)
  require(rgeos)
  set.seed(seed)
  
  # ... make coords
  # ... --- raster
  base.r <- raster(extent(shape), res = res)
  shape$is.shape <- 1
  base.r <- raster::rasterize(shape, base.r, field = "is.shape", background = NA)
  
  # ... --- points
  points <- na.omit(as.data.frame(base.r, xy = T))
  coords <- points[,1:2]
  
  # ... how many clusters?
  if(size.km2){
    require(geosphere)
    N <- max(c(1,round((areaPolygon(shape)/1000^2) / (size))))
  } else {
    N <- max(c(1,round(gArea(shape) / (size))))
  }
  
  
  # ... return if N == 1
  if(N == 1){
    tess <- shape
    tess@data <- data.frame(cell.id = 1)
    return(tess)
  }
  # ... regular sampling of centers
  set.seed(seed)
  m <- 0
  while(T & m < iter.max){
    m = m + 1
    centers <- spsample(shape, n = N, type = sample.type)@coords
    if(nrow(centers) == N){
      break
    }
  }
  
  # ... kmeans of coords
  set.seed(seed)
  km <- stats::kmeans(x = coords, centers = centers, algorithm = "Lloyd", iter.max = iter.max)
  
  # ... tesselate
  tess <- try(voronoi_poly(SpatialPoints(na.omit(km$centers)),
                           extent = extent(shape)))
  if(class(tess) == "try-error"){print(paste("ERROR IN VORONOI TESSELATION"))}
  
  # ... crop
  if(!is.null(shape)){
    if(length(shape) > 1){
      shape <- gUnaryUnion(shape, id = rep(1, length(shape)))
    }
    proj4string(shape) <- proj4string(tess)
    tess <- raster::crop(tess, shape)
  }
  
  # Return
  return(tess)
}
